#Data : 2020-5-19
#Author : Fengyuan Zhang (Franklin)
#Email : franklinzhang@foxmail.com
#Description : 

import os
import sys
import gdal
import ogr
import datetime
from osgeo import gdal, gdalconst
from pyproj import CRS

class CalculateLandcover:
    def __init__(self, landcover: str = "", sites: str = "", epsg: str = "epsg:3857", radius: float = 5000.0, nanvalue: int = -99998):
        self._landcover = landcover
        self._sites = sites
        self._prj = epsg
        self._radius = radius
        self._types = []
        self._nvalue = nanvalue

    #! private functions

    def _calculateTypes(self) -> int:
        if self._landcover == "":
            print("No land cover data, please set land cover data first!")
            return -1
        if not os.path.exists(self._landcover):
            print("Land cover data does not exist, please check it : " + self._landcover)
            return -1

        dataset = gdal.Open(self._landcover)
        if dataset == None:
            print("Can not open land cover data : " + self._landcover)
            return -1
        tif_width = dataset.RasterXSize
        tif_height = dataset.RasterYSize
        tif_bandsCount = dataset.RasterCount

        if tif_bandsCount > 2:
            print("Only support one band data : ")
            return -1

        tif_data = dataset.ReadAsArray(0, 0, tif_width, tif_height)
        tif_proj = dataset.GetProjection()

        for line in tif_data:
            for value in line:
                if value not in self._types:
                    self._types.append(value)
        return 1

    #! public functions
    def getLandcover(self) -> str:
        return self._landcover
    
    def setLandcover(self, landcover: str) -> None:
        self._landcover = landcover

    def getSites(self) -> str:
        return self._sites
    
    def setSites(self, sites: str) -> None:
        self._sites = sites

    def getPrj(self) -> str:
        return self._prj
    
    def setPrj(self, epsg: str) -> None:
        self._prj = epsg

    def getRadius(self) -> float:
        return self._radius

    def setRadius(self, radius: float) -> None:
        self._radius = radius

    def getTypes(self) -> list:
        return self._types.copy()

    def calculateByInversedDistance(self, log: bool = False) -> int:
        if self._calculateTypes() < 0:
            return -1

        dataset = gdal.Open(self._landcover)

        tif_geotrans = dataset.GetGeoTransform()
        tif_width = dataset.RasterXSize
        tif_height = dataset.RasterYSize
        tif_data = dataset.ReadAsArray(0, 0, tif_width, tif_height)

        pixelWidth = tif_geotrans[1]
        pixelHeight = tif_geotrans[5]
        originX = tif_geotrans[0] + pixelWidth / 2
        originY = tif_geotrans[3] + pixelHeight / 2

        geod = CRS(self._prj).get_geod()

        #! calculate cell area
        az12_wgs, az21_wgs, dist_wgs = geod.inv(originX, originY, originX + pixelWidth, originY)
        cellLen = dist_wgs
        area = cellLen * cellLen

        site_features = []
        site_geo = []

        mb_v = ogr.Open(self._sites, 1)
        mb_l = mb_v.GetLayer()
        f_source = mb_l.GetNextFeature()

        while f_source is not None:
            geom = f_source.GetGeometryRef()

            site_features.append(f_source)
            siteInfo = {"OID" : f_source.GetFID(), "location" : [geom.GetX(), geom.GetY()], "count" : 0, "rate" : 1}
            for typeItem in self._types:
                siteInfo[str(typeItem)] = 0.0
            site_geo.append(siteInfo)

            f_source = mb_l.GetNextFeature()

        for index_y in range(tif_height):
            for index_x in range(tif_width):
                strType = str(tif_data[index_y, index_x])
                if self._nvalue != -99998 and tif_data[index_y, index_x] == self._nvalue:
                    continue
                lons2 = originX + index_x * pixelWidth
                lats2 = originY + index_y * pixelHeight
                for site in site_geo:
                    az12_wgs, az21_wgs, dist_wgs = geod.inv(site["location"][0], site["location"][1], lons2, lats2)
                    if dist_wgs < self._radius:
                        site[strType] = site[strType] + (area*0.01)/(dist_wgs + 1)
                        site["count"] = site["count"] + 1
            if log:
                print(str(index_y + 1) + "/" + str(tif_height) + " Finished!")

        #! calculate the rate
        cell_max = 0
        for geo in site_geo:
            if geo["count"] > cell_max:
                cell_max = geo["count"]
                
        #! normalization
        for index in range(len(site_geo)):
            if site_geo[index]["count"] > cell_max/3:
                site_geo[index]["rate"] = site_geo[index]["count"]*1.000000/cell_max

        for index in range(len(site_geo)):
            for typeItem in self._types:
                site_geo[index][str(typeItem)] = round(site_geo[index][str(typeItem)]/site_geo[index]["rate"], 6)

        for typeItem in self._types:
            fieldType = ogr.FieldDefn("LC_" + str(typeItem), ogr.OFTReal)
            fieldType.SetWidth(64)
            fieldType.SetPrecision(6)
            mb_l.CreateField(fieldType)

        mb_l.ResetReading()
        f_source = mb_l.GetNextFeature()
        while f_source is not None:
            for geo in site_geo:
                if f_source.GetFID() == geo["OID"]:
                    for typeItem in self._types:
                        f_source.SetField("LC_" + str(typeItem), geo[str(typeItem)])
                        mb_l.SetFeature(f_source)
            f_source = mb_l.GetNextFeature()

        return 1

    def calculateByCount(self, log: bool = True) -> int:
        if self._calculateTypes() < 0:
            return -1

        dataset = gdal.Open(self._landcover)

        tif_geotrans = dataset.GetGeoTransform()
        tif_width = dataset.RasterXSize
        tif_height = dataset.RasterYSize
        tif_data = dataset.ReadAsArray(0, 0, tif_width, tif_height)

        pixelWidth = tif_geotrans[1]
        pixelHeight = tif_geotrans[5]
        originX = tif_geotrans[0] + pixelWidth / 2
        originY = tif_geotrans[3] + pixelHeight / 2

        geod = CRS(self._prj).get_geod()

        site_features = []
        site_geo = []

        mb_v = ogr.Open(self._sites, 1)
        mb_l = mb_v.GetLayer()
        f_source = mb_l.GetNextFeature()

        while f_source is not None:
            geom = f_source.GetGeometryRef()

            site_features.append(f_source)
            siteInfo = {"OID" : f_source.GetFID(), "location" : [geom.GetX(), geom.GetY()]}
            for typeItem in self._types:
                siteInfo[str(typeItem)] = 0
            site_geo.append(siteInfo)

            f_source = mb_l.GetNextFeature()

        for index_y in range(tif_height):
            for index_x in range(tif_width):
                strType = str(tif_data[index_y, index_x])
                if self._nvalue != -99998 and tif_data[index_y, index_x] == self._nvalue:
                    continue
                lons2 = originX + index_x * pixelWidth
                lats2 = originY + index_y * pixelHeight
                for site in site_geo:
                    az12_wgs, az21_wgs, dist_wgs = geod.inv(site["location"][0], site["location"][1], lons2, lats2)
                    if dist_wgs < self._radius:
                        site[strType] = site[strType] + 1
            if log:
                print(str(index_y + 1) + "/" + str(tif_height) + " Finished!")

        for geo in site_geo:
            for typeItem in self._types:
                geo[str(typeItem)] = round(geo[str(typeItem)], 3)

        for typeItem in self._types:
            fieldType = ogr.FieldDefn("LC_" + str(typeItem), ogr.OFTInteger)
            fieldType.SetWidth(32)
            mb_l.CreateField(fieldType)

        mb_l.ResetReading()
        f_source = mb_l.GetNextFeature()
        while f_source is not None:
            for geo in site_geo:
                if f_source.GetFID() == geo["OID"]:
                    for typeItem in self._types:
                        f_source.SetField("LC_" + str(typeItem), geo[str(typeItem)])
                        mb_l.SetFeature(f_source)
            f_source = mb_l.GetNextFeature()

        return 1